# 完全二叉树的节点个数： https://leetcode-cn.com/problems/count-complete-tree-nodes/submissions/

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

# O(n) 时间复杂度实现
class Solution:
    def countNodes(self, root: TreeNode) -> int:
        """
            后序遍历，左右根
        """
        def getNum(node):
            # 3. 递归终止边界条件： 当前节点为 None， 表示返回 0，按照题意
            if not node: return 0

            # 4. 递归求解 左右节点个数
            leftNum = getNum(node.left)
            rightNum = getNum(node.right)
        
            # 6. 返回当前整棵树的节点个数， + 1 是算上此时的根节点。它们整个是上一级回溯的子节点个数
            return leftNum + rightNum + 1 
        
        return getNum(root)




# 还可以采用二分查找法加 位运算来做~  确实有点难度 和难想
class Solution:
    def countNodes(self, root: TreeNode) -> int:
        """
            这道题使用深度优先，广度优先遍历都可以求。但是要求设计比O(n)遍历更快的算法
            但既是完全二叉树，这让我想到了二叉树的性质求
            层数从 0开始：
            1. k层的二叉树，最多有 2 ^ k  - 1 个结点（满二叉树）。
            2. 第 i 层上最多有 2 ^ (i - 1) 个结点
            3. 我们知道 k - 1层是满二叉树，很好求出节点数。 那么只需要确定第 k 层节点数， 并且第 k层节点的范围也是知道的，根据性质2.
            4. 范围有了，利用二分查找，确定第i层有几个节点。 用二进制表示， 举例子， 第五层最多有 16个结点， 第8个结点为（二进制）： 10111， 十进制表示是 第23个结点。
            第一个1代表第五层（第一个1代表层数）， 后面四个 0代表左子树， 1代表右子树， 从根节点推过来。 另外代表在本层的位置，下标从 0开始
            根节点 -> 左节点 -> 右节点 -> 右节点 -> 右节点（为当前节点）

            二分查找时间复杂度 O(logn), 查找该节点是否存在 O(logn)树的深度。  合在一起时间复杂度就是 O(logn * logn)

            还需要使用位运算来表示
        """
        # 定义判断节点是否存在的函数
        def exist(root: TreeNode, level: int, mid: int):
            """
                param:  root: 树的根节点
                param:  level: 树的层数，深度
                param:  mid: 当前判断是否存在的 节点的 序号，下标从 0 开始
            """
            # 1. 先计算出掩码： 用来求出 2 进制每位的值
            mask = 1 << (level - 2)  # level = 4, 01000
            cur = root
            # 2. 进行移位，看是否存在该节点 且移位次数， mask 大于0 控制
            while cur and mask > 0:
                # 与运算， 0 左， 1 右节点
                if mask & mid == 0:
                    cur = cur.left
                else:
                    cur = cur.right

                # 更新mask，右移下一位， 看下一步取左节点，右节点
                mask = mask >> 1

            return cur != None

        if not root: return 0

        # （1） 先找到树的深度, 一直找左节点， 一定是树的深度，对于完全二叉树。 注意数的深度是从 1开始的！！
        cur = root
        k = 0
        while cur:
            k += 1
            cur = cur.left
        # (2) 进行二分查找
        left, right = 2 ** (k-1), 2 ** k - 1
        while left < right:
            mid = (left + right + 1) // 2
            if exist(root, k, mid):
                left = mid
            else:
                right = mid - 1

        # left 的值，就是节点的个数
        return left
